package gov.nasa.jpf.symbc.modelcounter.sampler;

import java.util.BitSet;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;

import org.apfloat.Apfloat;

import gov.nasa.jpf.Config;
import gov.nasa.jpf.symbc.modelcounter.sampler.EvaluationVisitor.Result;
import gov.nasa.jpf.symbc.numeric.Constraint;
import gov.nasa.jpf.symbc.numeric.PathCondition;

/**
 * A model counter that "counts" models by sampling the variable space.
 */
public class ConfIntervalSampler extends Sampler {

	/**
	 * How many samples are taken between calculations of the SEM (standard error of the mean).
	 */
	private static final int SEM_FREQUENCY = 50;

	/**
	 * How many samples do we store before we decide to flush the entire tree?
	 */
	private static final int FLUSH_LIMIT = 4 * 1000 * 1000;

	/**
	 * The PC tree records information about path conditions and specifically
	 * the samples that have been taken in the past to evaluate them. Each node
	 * represents a constraint with zero or more children. A path condition is a
	 * path in the tree from a root to a node.
	 */
	protected static class PCTree {

		/**
		 * Total number of samples stored in all trees.
		 */
		private static int combinedCount = 0;

		/**
		 * The number of the constraint that splits the tree (into its true
		 * branches and false branches) at the current node. This number is
		 * always even since only "positive" constraints are stored in the tree.
		 */
		protected int constraintNumber;

		/**
		 * Children of the tree node that represents the case where the
		 * constraint is true.
		 */
		protected Map<Integer, PCTree> leftChildren;

		/**
		 * Children of the tree node that represents the case where the
		 * constraint is false.
		 */
		protected Map<Integer, PCTree> rightChildren;

		/**
		 * Total number of samples taken.
		 */
		protected int totalCount;

		/**
		 * Number of samples that evaluated to true.
		 */
		protected int trueCount;

		/**
		 * Number of samples that evaluated to false.
		 */
		protected int falseCount;

		/**
		 * The samples that evaluated to true.
		 */
		protected Set<Assignment> trueSamples;

		/**
		 * The samples that evaluated to false.
		 */
		protected Set<Assignment> falseSamples;

		/**
		 * The depth of this node. Not really needed except for debugging.
		 */
		protected int depth = 0;

		/**
		 * Constructor that initialize the tree node.
		 */
		public PCTree(int constraintNumber, PCTree parent) {
			this.constraintNumber = constraintNumber;
			leftChildren = new HashMap<Integer, PCTree>();
			rightChildren = new HashMap<Integer, PCTree>();
			totalCount = 0;
			trueCount = 0;
			falseCount = 0;
			trueSamples = new HashSet<Assignment>();
			falseSamples = new HashSet<Assignment>();
			depth = (parent != null) ? parent.depth + 1 : 0;
		}

		/**
		 * Return the number of samples that evaluated to true.
		 * 
		 * @return number of true samples
		 */
		public int getTrueCount() {
			return trueCount;
		}

		/**
		 * Increment the number of true samples *AS WELL AS* the total number of
		 * samples. Add the given assignment to the repository of samples.
		 * 
		 * @param assignment
		 *            the sample that evaluated as false
		 */
		public void addTrueSample(Assignment assignment) {
			trueSamples.add(assignment);
			trueCount++;
			totalCount++;
			combinedCount++;
		}

		/**
		 * Return the number of samples that evaluated to false.
		 * 
		 * @return number of false samples
		 */
		public int getFalseCount() {
			return falseCount;
		}

		/**
		 * Increment the number of false samples *AS WELL AS* the total number
		 * of samples. Add the given assignment to the repository of samples.
		 * 
		 * @param assignment
		 *            the sample that evaluated as false
		 */
		public void addFalseSample(Assignment assignment) {
			falseSamples.add(assignment);
			falseCount++;
			totalCount++;
			combinedCount++;
		}

		/**
		 * Return the total number of samples.
		 */
		public int getTotalCount() {
			return totalCount;
		}

		/**
		 * Return the combined count of all samples stored in a tree, somewhere.
		 */
		public int getCombinedCount() {
			return combinedCount;
		}

		/*
		 * (non-Javadoc)
		 * 
		 * @see java.lang.Object#toString()
		 */
		@Override
		public String toString() {
			StringBuffer b = new StringBuffer();
			b.append("[C:");
			b.append(constraintNumber);
			b.append(" T:");
			b.append(trueCount);
			b.append(" F:");
			b.append(falseCount);
			b.append(" D:");
			b.append(depth);
			b.append(']');
			return b.toString();
		}

	}

	/**
	 * The size of the 95% confidence interval to calculate. For example, if
	 * this value is 0.01, then the sampler would produce values such as [0.445
	 * .. 0.505].
	 */
	protected double confInterval = 0.01;

	/**
	 * Map from constraints to internal numbers.
	 */
	protected Map<Constraint, Integer> constraintLookupTable;

	/**
	 * Map from internal numbers to constraints.
	 */
	protected Map<Integer, Constraint> reverseLookupTable;

	/**
	 * Map from internal numbers (e.g., 132 which may represent x == y) to
	 * information about the samples taken to evaluate the probability of the
	 * constraint. This is needed because there may in fact be several trees
	 * (i.e., a forest) instead of just one.
	 */
	protected Map<Integer, PCTree> treeRoots;

	/**
	 * A visitor for constraints/expressions used to determine whether samples
	 * are true or false.
	 */
	protected EvaluationVisitor evaluator;

	/**
	 * Constructor for the sampler. Call the super-constructor and initialize
	 * the local fields.
	 * 
	 * @param config
	 *            a JPF configuration
	 */
	public ConfIntervalSampler(Config config) {
		super(config);
		double z = config.getDouble("symbolic.prob_conf_interval",
				Double.MAX_VALUE);
		if (z != Double.MAX_VALUE) {
			confInterval = z;
		}
		System.out.println("symbolic.prob_conf_interval=" + confInterval);
		constraintLookupTable = new HashMap<Constraint, Integer>();
		reverseLookupTable = new HashMap<Integer, Constraint>();
		treeRoots = new HashMap<Integer, PCTree>();
		evaluator = new EvaluationVisitor(null);
	}

	/**
	 * Return the probability that the given path condition is true.
	 * 
	 * Note: we assume that the constraint pathCondition.header is the most
	 * recent constraint added to the path condition.
	 * 
	 * @param pathCondition
	 *            the path condition to evaluate
	 * @return the probability that the path condition is true
	 * @see gov.nasa.jpf.symbc.modelcounter.ModelCounter#getConditionalProbability(gov.nasa.jpf.symbc.numeric.PathCondition,
	 *      double)
	 */
	@Override
	public Apfloat getConditionalProbability(PathCondition pathCondition,
			Apfloat parentProb) {
		int freshConstraint = lookup(pathCondition.header);
		boolean negated = (freshConstraint & 1) == 1;
		BitSet constraintSet = convertToBitSet(pathCondition.header.and);
//		System.out.println(">> PC:"
//				+ pathCondition.header.toString().replaceAll("\n", " "));
		// System.out.println(">> freshConstraint:" + freshConstraint);
		// System.out.println(">> negated:" + negated);
		// System.out.println(">> constraintSet:" + constraintSet);
		PCTree node = findPCTreeNode(freshConstraint & ~1, constraintSet);
		// System.out.println(">> node:" + node);
		double ci = 0.5, m = 0.0;
		boolean completed = false;
		for (int i = 0; i < sampleLimit; i+= SEM_FREQUENCY) {
			for (int j = 0; j < SEM_FREQUENCY; j++) {
				sample(pathCondition.header, node, negated);
			}
			double t = node.getTotalCount();
			if (t > 5) {
				double tc = node.getTrueCount();
				m = tc / t;
				double d = tc * m * m + (t - tc) * (1 - m) * (1 - m);
				ci = Math.sqrt(d / (t * t)) * 1.96;
				if (ci * 2.0 < confInterval) {
					completed = true;
					break;
				}
			}
		}
		Apfloat pr = null;
		if (negated) {
			pr = Apfloat.ONE.multiply(new Apfloat(node.getFalseCount())).divide(new Apfloat(node.getTotalCount()));
		} else {
			pr = Apfloat.ONE.multiply(new Apfloat(node.getTrueCount())).divide(new Apfloat(node.getTotalCount()));
		}
		Apfloat lo = pr.subtract(new Apfloat(ci));
		Apfloat hi = pr.add(new Apfloat(ci));
		System.out.println(">> [" + lo + " .. " + hi + "]"
				+ (completed ? "*" : "") + " PC:"
				+ pathCondition.header.toString().replaceAll("\n", " "));
//		System.out.println("%%%% " + node.getCombinedCount());
		if (node.getCombinedCount() > FLUSH_LIMIT) {
			treeRoots.clear();
		}
		return pr;
	}

	/**
	 * Returns an internal number that represents the given constraint. The
	 * routine handles negation is a special way. Suppose we have seen
	 * constraint X before, and assigned it the value 20. If this routine is
	 * subsequently called with !X, it returns 21.
	 * 
	 * @param constraint
	 *            a given constraint
	 * @return a unique number to identify the constraint
	 */
	private int lookup(Constraint constraint) {
		if (constraintLookupTable.containsKey(constraint)) {
			return constraintLookupTable.get(constraint);
		}
		// Try the negation of the constraint
		Constraint n = constraint.not();
		if (constraintLookupTable.containsKey(n)) {
			return constraintLookupTable.get(n) + 1;
		}
		// OK, we have to add the constraint to the table
		int v = constraintLookupTable.size() * 2;
		constraintLookupTable.put(constraint, v);
		reverseLookupTable.put(v, constraint);
		// System.out.println(">> ADD LOOKUP " + v + " -> " +
		// constraint.toString().replaceAll("\n", " "));
		return v;
	}

	/**
	 * Returns the constraint that corresponds to the internal number. This
	 * routine handles negation in accordance with the behaviour of
	 * {@link ConfIntervalSampler#lookup(Constraint)}.
	 * 
	 * @param constraintNumber
	 *            the number of the constraint to find
	 * @return the constraint that corresponds to the number
	 */
	private Constraint reverseLookup(int constraintNumber) {
		int p = constraintNumber & ~1;
		assert reverseLookupTable.containsKey(p) : "Unexpected, unknown constraint"
				+ constraintNumber;
		if ((p & 1) != 0) {
			return reverseLookupTable.get(p).not();
		} else {
			return reverseLookupTable.get(p);
		}
	}

	/**
	 * Convert a set of constraints (linked through their ".and" fields) to a
	 * bitset of their internal numbers.
	 * 
	 * @param constraintList
	 *            the linked list of constraints
	 * @return a bitset with the constraint numbers
	 */
	private BitSet convertToBitSet(Constraint constraintList) {
		BitSet b = new BitSet();
		while (constraintList != null) {
			b.set(lookup(constraintList));
			constraintList = constraintList.and;
		}
		return b;
	}

	/**
	 * Find (or create, if necessary) the node in the PC tree that corresponds
	 * to the path condition that is formed by the conjunction of a set of
	 * constraints. The constraints are identified by their internal numbers and
	 * the set is split into two parts: the (internal number of) the first
	 * constraint, and a bitset with (the internal numbers of) the rest of the
	 * constraints. This makes it easy to find the "parent" node of the node we
	 * are really looking for.
	 * 
	 * @param first
	 *            the internal number of the first constraint
	 * @param rest
	 *            a bitset with the internal numbers of the other constraints in
	 *            the conjunction
	 * @return the (perhaps newly created node) that represents the conjunction
	 */
	private PCTree findPCTreeNode(int first, BitSet rest) {
		// System.out.println(">> >> first:" + first + " rest:" + rest);
		if (rest.isEmpty()) {
			// We expect first (or its negation) to be a root
			if (treeRoots.containsKey(first)) {
				return treeRoots.get(first);
			}
			// Otherwise, try the negation (denoted by first + 1)
			if (treeRoots.containsKey(first + 1)) {
				return treeRoots.get(first + 1);
			}
			// Otherwise, create a new root for the tree
			// System.out.println(">> >> NEW ROOT");
			PCTree n = new PCTree(first, null);
			treeRoots.put(first, n);
			return n;
		}

		// Split the current rest into its own "first" and "rest" and find the
		// corresponding PC tree node
		int f = rest.length() - 1; // next "first"
		int g = f & ~1; // positive constraint number of f
		rest.clear(f); // remove f from rest
		PCTree p = findPCTreeNode(g, rest); // find the parent node

		// System.out.println(">> >> p:" + p);
		// Find/construct the node for (first, rest)
		PCTree n = null;
		if (p.constraintNumber == f) {
			// first is a "true" descendant of its parent
			n = p.leftChildren.get(first);
			if (n == null) {
				n = new PCTree(first, p);
				p.leftChildren.put(first, n);
				// Re-use previous samples, if any
				if ((first & 1) != 0) {
					if (!p.trueSamples.isEmpty()) {
						Constraint c = reverseLookup(first);
						for (Assignment a : p.trueSamples) {
							evaluator.reset(a);
							c.accept(evaluator);
							Result e = evaluator.getResult();
							if (e == Result.TRUE) {
								n.addFalseSample(a);
							} else if (e == Result.FALSE) {
								n.addTrueSample(a);
							}
						}
					}
				} else {
					if (!p.trueSamples.isEmpty()) {
						Constraint c = reverseLookup(first);
						for (Assignment a : p.trueSamples) {
							evaluator.reset(a);
							c.accept(evaluator);
							Result e = evaluator.getResult();
							if (e == Result.TRUE) {
								n.addTrueSample(a);
							} else if (e == Result.FALSE) {
								n.addFalseSample(a);
							}
						}
					}
				}
			}
		} else {
			// first is a "false" descendant of its parent
			n = p.rightChildren.get(first);
			if (n == null) {
				n = new PCTree(first, p);
				p.rightChildren.put(first, n);
				// Re-use previous samples, if any
				if ((first & 1) != 0) {
					if (!p.falseSamples.isEmpty()) {
						Constraint c = reverseLookup(first);
						for (Assignment a : p.falseSamples) {
							evaluator.reset(a);
							c.accept(evaluator);
							Result e = evaluator.getResult();
							if (e == Result.TRUE) {
								n.addFalseSample(a);
							} else if (e == Result.FALSE) {
								n.addTrueSample(a);
							}
						}
					}
				} else {
					if (!p.falseSamples.isEmpty()) {
						Constraint c = reverseLookup(first);
						for (Assignment a : p.falseSamples) {
							evaluator.reset(a);
							c.accept(evaluator);
							Result e = evaluator.getResult();
							if (e == Result.TRUE) {
								n.addTrueSample(a);
							} else if (e == Result.FALSE) {
								n.addFalseSample(a);
							}
						}
					}
				}
			}
		}

		return n;
	}

	/**
	 * Take one sample, evaluate the given conjunction of constraints with the
	 * sample, and update the tree node with the result. A sample is called
	 * "viable" if it evaluates as true for all by the first constraint in the
	 * conjunction. Only viable samples ultimately counted. This is how we
	 * implement condition probability.
	 * 
	 * @param constraints
	 *            the conjunction of constraints to evaluate
	 * @param node
	 *            node in the tree where information about the outcome of the
	 *            evaluation is stored
	 * @param negated
	 *            whether (or not) the node was originally constructed (and
	 *            really represents) the negation of first constrain in the
	 *            list; this is important for classifying the sample taken
	 */
	private void sample(Constraint constraints, PCTree node, boolean negated) {
		sampleCount++;
		if (negated) {
			Assignment a = new MapAssignment(this);
			Constraint c = constraints.and;
			while (c != null) {
				evaluator.reset(a);
				c.accept(evaluator);
				if (evaluator.getResult() != Result.TRUE) {
					break;
				}
				c = c.and;
			}
			if (c == null) {
				// Sample a is viable, so evaluate the header.
				evaluator.reset(a);
				constraints.accept(evaluator);
				Result e = evaluator.getResult();
				if (e == Result.TRUE) {
					node.addFalseSample(a);
				} else if (e == Result.FALSE) {
					node.addTrueSample(a);
				} // else e == Result.INVALID
			}
			// else { sample is not viable }
		} else {
			Assignment a = new MapAssignment(this);
			Constraint c = constraints.and;
			while (c != null) {
				evaluator.reset(a);
				c.accept(evaluator);
				if (evaluator.getResult() != Result.TRUE) {
					break;
				}
				c = c.and;
			}
			if (c == null) {
				// Sample a is viable, so evaluate the header.
				evaluator.reset(a);
				constraints.accept(evaluator);
				Result e = evaluator.getResult();
				if (e == Result.TRUE) {
					node.addTrueSample(a);
				} else if (e == Result.FALSE) {
					node.addFalseSample(a);
				} // else e == Result.INVALID
			}
			// else { sample is not viable }
		}
	}

	@Override
	public Apfloat getProbability(
			Set<PathCondition> pathConditions) {
		// TODO Auto-generated method stub
		return null;
	}

}
